"""
张量计算
"""

import torch


def test_add():
    print("""==========add========""")
    a = torch.rand(2, 3)
    b = torch.rand(2, 3)

    print("a")
    print(a)
    print("b")
    print(b)

    print("a+b")
    print(a + b)
    print("a.add(b)")
    print(a.add(b))
    print("torch.add(a,b)")
    print(torch.add(a, b))
    print("a.add_(b) 此方法会改变a的值")
    print(a.add_(b))  # 会返回结果
    print(a)

    print()


def test_sub():
    print("""==========sub========""")
    a = torch.rand(2, 3)
    b = torch.rand(2, 3)

    print("a")
    print(a)
    print("b")
    print(b)

    print("a-b")
    print(a - b)
    print("a.sub(b)")
    print(a.sub(b))
    print("torch.sub(a,b)")
    print(torch.sub(a, b))
    print("a.sub_(b) 此方法会改变a的值")
    print(a.sub_(b))  # 会返回结果
    print(a)

    print()


def test_mul():
    print("""==========哈达玛乘法（对应元素乘法）========""")
    a = torch.rand(2, 3)
    b = torch.rand(2, 3)

    print("a")
    print(a)
    print("b")
    print(b)

    print("a*b")
    print(a * b)
    print("a.mul(b)")
    print(a.mul(b))
    print("torch.mul(a,b)")
    print(torch.mul(a, b))
    print("a.mul_(b) 此方法会改变a的值")
    print(a.mul_(b))  # 会返回结果
    print(a)

    print()


def test_div():
    print("""==========div========""")
    a = torch.rand(2, 3)
    b = torch.rand(2, 3)

    print("a")
    print(a)
    print("b")
    print(b)

    print("a/b")
    print(a / b)
    print("a.div(b)")
    print(a.div(b))
    print("torch.div(a,b)")
    print(torch.div(a, b))
    print("a.div_(b) 此方法会改变a的值")
    print(a.div_(b))  # 会返回结果
    print(a)

    print()


def test_matmul():
    print("""==========二维矩阵乘法========""")
    a = torch.ones(2, 1)
    b = torch.ones(1, 2)

    print("a")
    print(a)
    print("b")
    print(b)

    print("a@b")
    print(a @ b)
    print("a.mm(b)")
    print(a.mm(b))
    print("a.matmul(b)")
    print(a.matmul(b))
    print("torch.matmul(a,b)")
    print(torch.matmul(a, b))
    print("torch.mm(a,b)")
    print(torch.mm(a, b))
    print()


def test_high_matmul():
    print("""==========高位矩阵乘法(实际是后两维的矩阵乘法，前面的维度一致）(不支持mm算子)========""")
    a = torch.ones(1, 2, 3, 4)
    b = torch.ones(1, 2, 4, 3)

    print("a")
    print(a)
    print("b")
    print(b)

    print("a@b")
    print(a @ b)

    print("a.matmul(b)")
    print(a.matmul(b))

    print("torch.matmul(a,b)")
    print(torch.matmul(a, b))

    print("a.matmul(b).shape")
    print(a.matmul(b).shape)
    print()


def test_pow():
    print("""==========pow========""")
    a = torch.tensor([1, 2])

    print("a")
    print(a)

    print("a**3")
    print(a ** 3)
    print("a.pow(3)")
    print(a.pow(3))
    print("torch.pow(a,3)")
    print(torch.pow(a, 3))
    print("a.pow_(3) 此方法会改变a的值")
    print(a.pow_(3))  # 会返回结果
    print(a)

    print()


def test_exp():
    print("""==========exp========""")
    # exp_会修改原始的tensor张量，e是float,因此需要原始张量的类型也要是float
    a = torch.tensor([1, 2], dtype=torch.float32)

    print("a")
    print(a)

    print("a.exp()")
    print(a.exp())

    print("a.exp_()")
    print(a.exp_())

    a = torch.tensor([1, 2], dtype=torch.float32)
    print("torch.exp(a)")
    print(torch.exp(a))

    print("torch.exp_(a)")
    print(torch.exp_(a))
    print()


def test_log():
    print("""==========log  log2  log10========""")
    # log_会修改原始的tensor张量，e是float,因此需要原始张量的类型也要是float
    a = torch.tensor([10, 2], dtype=torch.float32)

    print("a")
    print(a)

    print("a.log()")
    print(a.log())

    print("a.log_()")
    print(a.log_())

    a = torch.tensor([10, 2], dtype=torch.float32)
    print("torch.log(a)")
    print(torch.log(a))

    print("torch.log_(a)")
    print(torch.log_(a))
    print()


def test_sqrt():
    print("""==========sqrt========""")
    # log_会修改原始的tensor张量，e是float,因此需要原始张量的类型也要是float
    a = torch.tensor([10, 2], dtype=torch.float32)

    print("a")
    print(a)

    print("a.sqrt()")
    print(a.sqrt())

    print("a.sqrt_()")
    print(a.sqrt_())

    a = torch.tensor([10, 2], dtype=torch.float32)
    print("torch.sqrt(a)")
    print(torch.sqrt(a))

    print("torch.sqrt_(a)")
    print(torch.sqrt_(a))
    print()


if __name__ == '__main__':
    # test_add()
    # test_sub()
    # test_mul()
    # test_div()
    # test_matmul()
    # test_high_matmul()
    # test_pow()
    # test_exp()
    # test_log()
    test_sqrt()
